/*
 * Copyright 2002-2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.naja.auth2clientfirst.support;

import org.redisson.api.RBucket;
import org.redisson.api.RedissonClient;
import org.springframework.boot.autoconfigure.web.ServerProperties;
import org.springframework.dao.DataRetrievalFailureException;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.client.OAuth2AuthorizedClient;
import org.springframework.security.oauth2.client.OAuth2AuthorizedClientService;
import org.springframework.security.oauth2.client.registration.ClientRegistration;
import org.springframework.security.oauth2.client.registration.ClientRegistrationRepository;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.stereotype.Service;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.time.Duration;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;

public class RedisOAuth2AuthorizedClientService implements OAuth2AuthorizedClientService {
    private ClientRegistrationRepository clientRegistrationRepository;
    private RedissonClient redisson;
    private final static Long TIMEOUT = TimeUnit.DAYS.toSeconds(7);
    private static final String AUTHORIZATION_CLIENT = "oauth2:client:token:%s:%s";

    public RedisOAuth2AuthorizedClientService(RedissonClient redisson,
                                              ClientRegistrationRepository clientRegistrationRepository) {
        this.redisson = redisson;
        this.clientRegistrationRepository = clientRegistrationRepository;
    }

    @Override
    public <T extends OAuth2AuthorizedClient> T loadAuthorizedClient(String clientRegistrationId,
                                                                     String principalName) {
        Assert.hasText(clientRegistrationId, "clientRegistrationId cannot be empty");
        Assert.hasText(principalName, "principalName cannot be empty");
        RBucket<Map<String, Object>> bucket = redisson.getBucket(String.format(AUTHORIZATION_CLIENT, clientRegistrationId, principalName));
        Map<String, Object> map = bucket.get();
        if (map == null) {
            return null;
        }

        ClientRegistration clientRegistration = this.clientRegistrationRepository.findByRegistrationId(clientRegistrationId);
        if (clientRegistration == null) {
            throw new DataRetrievalFailureException(
                    "The ClientRegistration with id '" + clientRegistrationId + "' exists in the data source, "
                            + "however, it was not found in the ClientRegistrationRepository.");
        }
        OAuth2AccessToken.TokenType tokenType = null;
        if (OAuth2AccessToken.TokenType.BEARER.getValue().equalsIgnoreCase((String) map.get("accessTokenType"))) {
            tokenType = OAuth2AccessToken.TokenType.BEARER;
        }
        String tokenValue = (String) map.get("accessTokenValue");
        Instant issuedAt = (Instant) map.get("accessTokenIssuedAt");
        Instant expiresAt = (Instant) map.get("accessTokenExpiresAt");
        Set<String> scopes = Collections.emptySet();
        String accessTokenScopes = (String) map.get("accessTokenScopes");
        if (accessTokenScopes != null) {
            scopes = StringUtils.commaDelimitedListToSet(accessTokenScopes);
        }
        OAuth2AccessToken accessToken = new OAuth2AccessToken(tokenType, tokenValue, issuedAt, expiresAt, scopes);
        OAuth2RefreshToken refreshToken = null;
        String refreshTokenValue = (String) map.get("refreshTokenValue");
        if (refreshTokenValue != null) {
            tokenValue = refreshTokenValue;
            issuedAt = (Instant) map.get("refreshTokenIssuedAt");
            expiresAt = (Instant) map.get("refreshTokenExpiresAt");
            refreshToken = new OAuth2RefreshToken(tokenValue, issuedAt, expiresAt);
        }
        return (T) new OAuth2AuthorizedClient(clientRegistration, principalName, accessToken, refreshToken);

    }

    @Override
    public void saveAuthorizedClient(OAuth2AuthorizedClient authorizedClient, Authentication principal) {
        Assert.notNull(authorizedClient, "authorizedClient cannot be null");
        Assert.notNull(principal, "principal cannot be null");

        ClientRegistration clientRegistration = authorizedClient.getClientRegistration();
        OAuth2AccessToken accessToken = authorizedClient.getAccessToken();
        OAuth2RefreshToken refreshToken = authorizedClient.getRefreshToken();

        long between = ChronoUnit.SECONDS.between(accessToken.getIssuedAt(), accessToken.getExpiresAt());
        Map<String, Object> map = new HashMap<>();
        map.put("clientRegistrationId", clientRegistration.getRegistrationId());
        map.put("principalName", principal.getName());
        map.put("accessTokenType", accessToken.getTokenType().getValue());
        map.put("accessTokenValue", accessToken.getTokenValue());
        map.put("accessTokenIssuedAt", accessToken.getIssuedAt());
        map.put("accessTokenExpiresAt", accessToken.getExpiresAt());
        String accessTokenScopes = null;
        if (!CollectionUtils.isEmpty(accessToken.getScopes())) {
            accessTokenScopes = StringUtils.collectionToDelimitedString(accessToken.getScopes(), ",");
        }
        map.put("accessTokenScopes", accessTokenScopes);
        if (refreshToken != null) {
            String refreshTokenValue = refreshToken.getTokenValue();
            map.put("refreshTokenValue", refreshTokenValue);
            map.put("refreshTokenIssuedAt", refreshToken.getIssuedAt());
            map.put("refreshTokenExpiresAt", refreshToken.getExpiresAt());
            if (refreshToken.getExpiresAt() != null && refreshToken.getExpiresAt().isAfter(accessToken.getExpiresAt())) {
                between = ChronoUnit.SECONDS.between(accessToken.getIssuedAt(), refreshToken.getExpiresAt());
            }
        }
        RBucket<Map<String, Object>> bucket = redisson.getBucket(String.format(AUTHORIZATION_CLIENT, clientRegistration.getRegistrationId(), principal.getName()));
        bucket.set(map, TIMEOUT, TimeUnit.SECONDS);
    }

    @Override
    public void removeAuthorizedClient(String clientRegistrationId, String principalName) {
        Assert.hasText(clientRegistrationId, "clientRegistrationId cannot be empty");
        Assert.hasText(principalName, "principalName cannot be empty");
        RBucket<Map<String, Object>> bucket = redisson.getBucket(String.format(AUTHORIZATION_CLIENT, clientRegistrationId, principalName));
        bucket.delete();
    }
}
